#!/usr/bin/env python
# -*- coding:utf-8 -*-
# author: Jiguang Peng
# created: 2019/6/27 19:13

import os
import re
import sys
import random
import string
from collections import namedtuple
from typing import Optional
from pyfaidx import Fasta
from .pvs1 import PVS1
from .cnv import PVS1CNV, CNVRecord
from .read_data import trans_gene, gene_trans, gene_alias
from .read_data import transcripts_hg19, transcripts_hg38
from .utils import get_transcript, vep_consequence_trans, VCFRecord


lof_type = ['frameshift', 'nonsense', 'splice-5', 'splice-3', 'init-loss']
vep_lof_list = ['frameshift', 'stop_gained', 'splice_donor', 'splice_acceptor', 'start_lost']
VAR = namedtuple('VAR', ('varid', 'gene', 'trans', 'canonical', 'pick', 'record'))


def id_generator(size=6, chars=string.ascii_uppercase + string.digits):
    return ''.join(random.choice(chars) for _ in range(size))


class AutoPVS1:
    """Run AutoPVS1"""

    def __init__(self, vcf_var: str, vep_outs: list, genome_version: str, genome: Fasta):
        if genome_version in ['hg19', 'GRCh37']:
            genome_version = 'hg19'
        elif genome_version in ['hg38', 'GRCh38']:
            genome_version = 'hg38'
        else:
            raise ValueError("Genome version must be hg19/GRCh37 or hg38/GRCh38.")
        vcfrecord = self.get_vcfrecord(vcf_var)
        vep_trans, consequence, hgvs_c, hgvs_p = self.vep_filter(vep_outs)
        transcript = get_transcript(vep_trans, transcripts_hg19) if genome_version == 'hg19' else get_transcript(vep_trans, transcripts_hg38)
        islof = consequence in lof_type
        self.pvs1: Optional[PVS1] = None
        if islof:
            self.pvs1 = PVS1(vcfrecord, consequence, hgvs_c, hgvs_p, transcript, genome_version, genome)

    @staticmethod
    def get_vcfrecord(vcf) -> VCFRecord:
        if re.match(r'^((?:chr)?[\dXYMT]{1,2})-(\d+)-([ATCG]+)-([ATCG]+)$', vcf, re.I):
            v = vcf.split("-")
            v[0] = re.sub('chr', '', v[0], flags=re.I)
            return VCFRecord(v[0], int(v[1]), v[2], v[3])
        else:
            raise TypeError("Wrong VCF Record")

    @staticmethod
    def vep_filter(vep_outs: list) -> tuple:
        var_annos = list()
        for vep_out in vep_outs:
            if vep_out['SYMBOL'] == '-':
                vep_out['SYMBOL'] = trans_gene.get(vep_out['Feature'], "NA")
            if vep_out['SYMBOL'] == 'NA':
                vep_out['SYMBOL'] = trans_gene.get(vep_out['Feature'].split(".")[0], "NA")
            if vep_out['SYMBOL'] in gene_alias:
                vep_out['SYMBOL'] = gene_alias.get(vep_out['SYMBOL'])
            var_anno = VAR(vep_out['Uploaded_variation'], vep_out['SYMBOL'],
                           vep_out['Feature'], vep_out['CANONICAL'], vep_out['PICK'], vep_out)
            var_annos.append(var_anno)

        final_choose = list(filter(lambda anno: gene_trans.get(anno.gene) == anno.trans or
                                   gene_trans.get(anno.gene, "na").split(".")[0] == anno.trans.split(".")[0], var_annos))
        final = ''
        if len(final_choose) > 1:
            for var_anno in final_choose:
                if var_anno.pick == '1':
                    final = var_anno
            if not final:
                final = final_choose[0]
        elif len(final_choose) == 1:
            final = final_choose[0]
        else:
            final = var_annos[0]
        vep_trans = final.record['Feature']
        vep_consequence = final.record['Consequence'].replace('_variant', '')
        consequence = vep_consequence_trans(vep_consequence)
        hgvs_c = final.record['HGVSc']
        hgvs_p = final.record['HGVSp'].replace('%3D', '=')
        return vep_trans, consequence, hgvs_c, hgvs_p


class AutoPVS1CNV:
    """Auto PVS1 for CNV"""

    def __init__(self, bed_var: str, vep_outs: list, genome_version: str):
        if genome_version in ['hg19', 'GRCh37']:
            genome_version = 'hg19'
        elif genome_version in ['hg38', 'GRCh38']:
            genome_version = 'hg38'
        else:
            raise ValueError("wrong genome version, use hg38/GRCh38 by default")
        cnvrecord = self.get_cnvrecord(bed_var)
        self.pvs1: Optional[PVS1CNV] = None
        if cnvrecord.type:
            consequence, vep_trans = self.vep_filter(vep_outs)
            transcript = get_transcript(vep_trans, transcripts_hg19) if genome_version == 'hg19' else get_transcript(vep_trans, transcripts_hg38)
            self.pvs1 = PVS1CNV(cnvrecord, consequence, transcript, genome_version)

    @staticmethod
    def get_cnvrecord(cnv) -> CNVRecord:
        cnvpattern1 = re.compile(r'(del|dup|tdup|ntdup)\((chr)?(\d+|X|Y):(\d+)-(\d+)\)', re.I)
        cnvpattern2 = re.compile(r'(chr)?(\d+|X|Y)-(\d+)-(\d+)-(del|dup|tdup|ntdup)', re.I)
        cnvmatch1 = cnvpattern1.match(cnv)
        cnvmatch2 = cnvpattern2.match(cnv)
        if cnvmatch1:
            return CNVRecord(cnvmatch1.group(3), int(cnvmatch1.group(4)), int(cnvmatch1.group(5)), cnvmatch1.group(1).upper())
        elif cnvmatch2:
            return CNVRecord(cnvmatch2.group(2), int(cnvmatch2.group(3)), int(cnvmatch2.group(4)), cnvmatch2.group(5).upper())
        else:
            return CNVRecord(chrom=None, start=None, end=None, type=None)

    @staticmethod
    def vep_filter(vep_outs: list) -> tuple:
        var_annos = list()
        for vep_out in vep_outs:
            if vep_out['SYMBOL'] == '-':
                vep_out['SYMBOL'] = trans_gene.get(vep_out['Feature'], "-")
            if vep_out['SYMBOL'] == '-':
                vep_out['SYMBOL'] = trans_gene.get(vep_out['Feature'].split(".")[0], "NA")
            var_anno = VAR(vep_out['Uploaded_variation'], vep_out['SYMBOL'],
                           vep_out['Feature'], vep_out['CANONICAL'], vep_out['PICK'], vep_out)
            var_annos.append(var_anno)
        final_choose = list(filter(lambda anno: gene_trans.get(anno.gene) == anno.trans, var_annos))
        if len(final_choose) == 0:
            final_choose = list(filter(lambda anno: gene_trans.get(anno.gene, "na").split(".")[0] ==
                                       anno.trans.split(".")[0], var_annos))
        final = ''
        if len(final_choose) > 1:
            for var_anno in final_choose:
                if var_anno.pick == '1':
                    final = var_anno
            if not final:
                final = final_choose[0]
        elif len(final_choose) == 1:
            final = final_choose[0]
        else:
            final = var_annos[0]
        consequence = final.record['Consequence'].replace('_variant', '')
        vep_trans = final.record['Feature']
        return consequence, vep_trans
